package trainer;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.HashMap;

/**
 * The trainer class
 * Usage: java trainer.Trainer trainA.txt modelA.dat
 * @author Yunkai Zhang
 * 
 */
public class Trainer {

  private String[] types = { "Unigram", "Bigram", "Trigram", "Trigram" };

  private int K = 100;

  private final int N = 3;

  private static double dl = 0.2;

  private final double E = 0.00001;

  private ArrayList<ArrayList<String>> folders;

  private HashMap<String, Double> finalModel = new HashMap<String, Double>();

  private double[] weight;

  private int modelType;

  private int numOfSentences = 0;

  public static void main(String[] args) {
    if (args.length < 2) {
      System.out.println("Usage: java trainer trainA.txt modelA.dat");
    } else {
      // String datasource = "trainA.txt";
      String datasource = args[0];
      String modelfile = args[1];
      // String modelfile = "modelA.dat";

      HeldOutTest hot = new HeldOutTest();
      dl = hot.HeldOutTest(datasource);

      Trainer trainer = new Trainer();

      trainer.Load(datasource);

      trainer.CrossValidation();

      trainer.GenerateModel();

      trainer.save2file(modelfile);
    }
  }

  /**
   * Trainer Constructor
   */
  public Trainer() {

    weight = new double[N];

    folders = new ArrayList<ArrayList<String>>();

    for (int i = 0; i < K; i++) {
      ArrayList<String> folder = new ArrayList<String>();
      folder.add("<S>");
      folder.add("<S>");
      folders.add(folder);
    }
  }

  /**
   * Save model to file
   * 
   * @param filename
   */
  public void save2file(String filename) {

    try {
      String output = filename;
      File op = new File(output);
      FileOutputStream fops;

      fops = new FileOutputStream(op);
      OutputStreamWriter osw = new OutputStreamWriter(fops);

      osw.write(types[modelType] + ":\n");

      for (String string : this.finalModel.keySet()) {
        osw.write(string + "\t" + this.finalModel.get(string) + "\n");
      }

      osw.close();
      fops.close();
    } catch (FileNotFoundException e) {
      e.printStackTrace();
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  /**
   * Generate the interpolated Model based on the trained weights
   */
  public void GenerateModel() {

    System.out.println("=============== Begin Generating Model ================");

    ArrayList<String> allData = new ArrayList<String>();
    for (int i = 0; i < K; i++) {
      allData.addAll(folders.get(i));
    }

    HashMap<String, Double> interpolated = new HashMap<String, Double>();

    ArrayList<HashMap<String, Double>> models;

    TrainModel tm = new TrainModel(allData, dl);

    models = tm.getProb();

    HashMap<String, Double> unigram = models.get(0);
    HashMap<String, Double> bigram = models.get(1);
    HashMap<String, Double> trigram = models.get(2);

    for (String string : trigram.keySet()) {
      String[] tags = string.split(" ");
      double prob = weight[2] * trigram.get(string);
      if (bigram.containsKey(tags[1] + " " + tags[2]))
        prob += weight[1] * bigram.get(tags[1] + " " + tags[2]);
      if (unigram.containsKey(tags[2]))
        prob += weight[0] * unigram.get(tags[2]);

      interpolated.put(string, prob);
    }

    models.add(interpolated);

    // compareModel(models, interpolated);
    finalModel = models.get(modelType);

    if (modelType == N)
      System.out.println("Best Model is Interpolated Model");
    else
      System.out.println("Best Model is " + types[modelType]);

    System.out.println("================ End Generating Model =================");
  }

  /**
   * Load data from training data
   * 
   * @param filename
   */
  private void Load(String filename) {
    try {
      BufferedReader in = new BufferedReader(new FileReader(filename));
      String line = "";

      while ((line = in.readLine()) != null) {
        if (line.equalsIgnoreCase("<PERIOD>")) {
          numOfSentences++;
        }
      }
      in.close();
      K = Math.min(K, numOfSentences); // impose ceiling of numOfSentences
      K = Math.max(K, 1); // impose floor of 1

      int i = 0;
      in = new BufferedReader(new FileReader(filename));
      while ((line = in.readLine()) != null) {
        folders.get(i).add(line);
        if (line.equalsIgnoreCase("<PERIOD>")) {
          folders.get(i).add("<S>");
          folders.get(i).add("<S>");
          i++;
          i = i % K;
        }
      }
      in.close();
    } catch (IOException e) {
    }
  }

  /**
   * Perform Cross Validation
   */
  public void CrossValidation() {

    int[] winnerCounts = new int[N + 1];

    double[][] weights = new double[K][3];

    Interpolater interpolater = new Interpolater(N, E);

    System.out.println("============== Cross Validation Begins ==============");

    if (K > 2) {
      for (int i = 0; i < K; i++) {
        ArrayList<String> trainingdata = new ArrayList<String>();
        ArrayList<String> heldoutdata = new ArrayList<String>();
        ArrayList<String> evaldata = new ArrayList<String>();
        ArrayList<HashMap<String, Double>> models;

        for (int j = 0; j < K; j++) {
          if (j == (i + 1) % K)
            evaldata.addAll(folders.get(j));
          else if (j == i)
            heldoutdata.addAll(folders.get(j));
          else
            trainingdata.addAll(folders.get(j));
        }

        System.out.println("==================== Folder No." + (i + 1) + " ====================");
        System.out.println("EM on folder No." + (i + 1) + " and evaluate on folder No:"
                + (((i + 1) % K) + 1));

        TrainModel tm = new TrainModel(trainingdata, dl);

        models = tm.getProb();

        weights[i] = interpolater.Interpolate(models, heldoutdata, evaldata, winnerCounts);

      }
    } else if (K == 2) {
      // Only two folers, not enough data for training, held-out set and evaluation
      for (int i = 0; i < 2; i++) {

        System.out.println("==================== Folder No." + (i + 1) + " ====================");
        System.out.println("EM on folder No." + (i + 1) + " and evaluate on folder No:"
                + (((i + 1) % K) + 1));

        TrainModel tm = new TrainModel(folders.get(i % K), dl);

        ArrayList<HashMap<String, Double>> models = tm.getProb();

        weights[0] = interpolater.Interpolate(models, folders.get((i + 1) % K), folders.get(i),
                winnerCounts);
      }
    } else {
      // Only one fold, which is not enough for training and held-out set
      TrainModel tm = new TrainModel(folders.get(0), dl);

      ArrayList<HashMap<String, Double>> models = tm.getProb();

      weights[0] = interpolater.Interpolate(models, folders.get(0), folders.get(0), winnerCounts);
    }

    System.out.println("============== Cross Validation Summary ==============");

    AvgWeight(weights);

    modelType = max(winnerCounts);

    System.out.println("============== Cross Validation Finishes ==============");

  }

  /**
   * Return the index of the element with the max value in the array.
   */
  private int max(int[] counts) {
    int index = 0;
    int maxValue = Integer.MIN_VALUE;
    for (int i = 0; i < counts.length; i++) {
      if (i == N)
        System.out.println("Interpolated Model Wins " + counts[i] + " times ");
      else
        System.out.println(types[i] + " Model Wins " + counts[i] + " times ");
      if (maxValue < counts[i]) {
        index = i;
        maxValue = counts[i];
      }
    }
    return index;
  }

  /**
   * Get the Average Weights
   * @param weights
   */
  public void AvgWeight(double[][] weights) {
    double[] sum = new double[N];
    for (int i = 0; i < K; i++) {
      for (int j = 0; j < N; j++) {
        sum[j] += weights[i][j];
      }
    }

    for (int j = 0; j < N; j++)
      weight[j] = sum[j] / K;

  }

}
